﻿/*
* Copyright (c) 2012 Kyle Hayward
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/
using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Drawing.Imaging;
using System.Linq;
using System.Text;
using System.IO;
using System.Windows.Forms;
using System.Runtime.InteropServices;

using OptixDotNet;
using OptixDotNet.Math;
using OptixDotNet.Utils;

using OptixDotNetUtils;

using Optix = OptixDotNet;

namespace Sample
{
    public struct BasicLight
    {
        public Vector3 Position;
        public Vector3 Color;
        public int CastsShadow;
        int padding;      // make this structure 32 bytes -- powers of two are your friend!
    }

    public class Tutorial : OptixForm
    {
        int mTutorial = 0;

        public Tutorial( int tutorial )
        {
            //Ensure that we have a valid Nvidia GPU
            Utils.DetectGPUCompatibility();

            mTutorial = tutorial;
        }

        protected override void Initialize()
        {
            string rayGenPath = Path.GetFullPath( "../../Assets/Cuda/tutorial" + mTutorial + ".ptx" );

            /*-----------------------------------------------
             * Create the Optix context
             *-----------------------------------------------*/
            Context = new Context();
            Context.RayTypeCount = 2;
            Context.EntryPointCount = 1;
            Context.StackSize = 4096;

            /*-----------------------------------------------
             * Create the ray-generation and exception programs
             *-----------------------------------------------*/
            Program rayGen = new Program( Context, rayGenPath, mTutorial < 11 ? "pinhole_camera" : "env_camera" );
            Program exception = new Program( Context, rayGenPath, "exception" );
            Program miss = new Program( Context, rayGenPath, mTutorial < 5 ? "miss" : "envmap_miss" );
            miss[ "bg_color" ].Set( 100 / 255.0f, 149 / 255.0f, 237 / 255.0f );
            exception[ "bad_color" ].SetFloat3( new Vector3( 1.0f, 0.0f, 0.0f ) );

            Context.SetRayGenerationProgram( 0, rayGen );
            Context.SetExceptionProgram( 0, exception );
            Context.SetRayMissProgram( 0, miss );

            /*-----------------------------------------------
             * Create lights
             *-----------------------------------------------*/
            BasicLight[] lights = new BasicLight[ 1 ];
            lights[ 0 ].Position = new Vector3( -5.0f, 60.0f, -16.0f );
            lights[ 0 ].Color = new Vector3( 1.0f, 1.0f, 1.0f );
            lights[ 0 ].CastsShadow = 1;

            BufferDesc desc = new BufferDesc() { Width = (uint)lights.Length, Format = Format.User, Type = BufferType.Input, 
                                                 ElemSize = (uint)Marshal.SizeOf( typeof( BasicLight ) ) };
            Optix.Buffer lightsBuffer = new Optix.Buffer( Context, desc );
            lightsBuffer.SetData<BasicLight>( lights );

            Context["lights"].Set( lightsBuffer );

            /*-----------------------------------------------
             * Create noise texture
             *-----------------------------------------------*/
            if ( mTutorial >= 8 )
            {
                uint noiseTexDim = 64;
                desc = new BufferDesc() { Width = noiseTexDim, Height = noiseTexDim, Depth = noiseTexDim, Format = Format.Float, Type = BufferType.Input };
                Optix.Buffer noiseBuffer = new Optix.Buffer( Context, desc );

                Random rand = new Random();
                BufferStream stream = noiseBuffer.Map();
                for ( int i = 0; i < noiseTexDim * noiseTexDim * noiseTexDim; i++ )
                {
                    stream.Write<float>( (float)rand.NextDouble() );
                }
                noiseBuffer.Unmap();

                TextureSampler noiseTex = new TextureSampler( Context, TextureSamplerDesc.GetDefault( WrapMode.Repeat ) );
                noiseTex.SetBuffer( noiseBuffer, 0 );

                Context[ "noise_texture" ].Set( noiseTex );
            }

            /*-----------------------------------------------
             * Load enivronment map texture
             *-----------------------------------------------*/
            LoadEnvMap();

            /*-----------------------------------------------
             * Load the geometry
             *-----------------------------------------------*/
            CreateGeometry();

            /*-----------------------------------------------
             * Create the output buffer
             *-----------------------------------------------*/
            CreateOutputBuffer( Format.UByte4 );
            Context[ "output_buffer" ].Set( OutputBuffer );

            /*-----------------------------------------------
            * Finally compile the optix context, and build the accel tree
            *-----------------------------------------------*/
            SetCamera();

            Context[ "max_depth" ].Set( 100 );
            Context[ "radiance_ray_type" ].Set( 0u );
            Context[ "shadow_ray_type" ].Set( 1u );
            Context[ "frame_number" ].Set( 0u );
            Context[ "scene_epsilon" ].Set( .001f );
            Context[ "importance_cutoff" ].Set( 0.01f );
            Context[ "ambient_light_color" ].Set( 0.31f, 0.33f, 0.28f );

            Context.Compile();
            Context.BuildAccelTree();

            //very loose calculation of number of rays
            float numSecondaryRays = 0;
            if ( mTutorial >= 9 )
                numSecondaryRays = 2.5f; //only convex hull casts refraction rays
            else if ( mTutorial >= 8 )
                numSecondaryRays = 2;
            else if ( mTutorial >= 4 )
                numSecondaryRays = 1.5f; //only the floor casts reflection rays, so give aproximate
            else if ( mTutorial >= 3 )
                numSecondaryRays = 1;

            RaysTracedPerFrame = (int)( Width * Height * ( numSecondaryRays + 1 ) );
        }

        protected override void RayTrace()
        {
            base.RayTrace();

            Context.Launch( 0, (uint)Width, (uint)Height );
        }

        private void CreateGeometry()
        {
            string shaderPath = Path.GetFullPath( "../../Assets/Cuda/tutorial" + mTutorial + ".ptx" );
            string boxPath = Path.GetFullPath( "../../Assets/Cuda/box.ptx" );
            string parrallelPath = Path.GetFullPath( "../../Assets/Cuda/parallelogram.ptx" );

            // Create box
            Geometry box = new Geometry( Context );
            box.PrimitiveCount = 1;
            box.BoundingBoxProgram = new Program( Context, boxPath, "box_bounds" ); ;
            box.IntersectionProgram = new Program( Context, boxPath, "box_intersect" ); ;
            box[ "boxmin" ].Set( -2.0f, 0.0f, -2.0f );
            box[ "boxmax" ].Set( 2.0f, 7.0f, 2.0f );

            Geometry chull = null;
            if ( mTutorial >= 9 )
            {
                chull = new Geometry( Context );
                chull.PrimitiveCount = 1u;
                chull.BoundingBoxProgram = new Program( Context, shaderPath, "chull_bounds" );
                chull.IntersectionProgram = new Program( Context, shaderPath, "chull_intersect" );

                uint nsides = 6;
                float radius = 1;
                Vector3 xlate = new Vector3( -1.4f, 0, -3.7f );

                BufferDesc desc = new BufferDesc() { Width = nsides + 2u, Type = BufferType.Input, Format = Format.Float4 };
                Optix.Buffer chullBuffer = new Optix.Buffer( Context, desc );

                float angle = 0.0f;
                BufferStream stream = chullBuffer.Map();
                for ( uint i = 0; i < nsides; i++ )
                {
                    angle = (float)i / (float)nsides * (float)Math.PI * 2.0f;
                    float x = (float)Math.Cos(angle);
                    float y = (float)Math.Sin(angle);

                    stream.Write<Vector4>( Vector4.CreatePlane( new Vector3( x, 0, y ), new Vector3( x * radius, 0, y * radius ) + xlate ) );
                }

                float min = 0.02f;
                float max = 3.5f;
                angle = 5.0f / (float)nsides * (float)Math.PI * 2.0f;
                stream.Write<Vector4>( Vector4.CreatePlane( new Vector3( 0, -1, 0 ), new Vector3( 0, min, 0 ) + xlate ) );
                stream.Write<Vector4>( Vector4.CreatePlane( new Vector3( (float)Math.Cos( angle ), 0.7f, (float)Math.Sin( angle ) ), 
                                                            new Vector3( 0, max, 0 ) + xlate ) );

                chullBuffer.Unmap();

                chull[ "planes" ].Set( chullBuffer );
                chull[ "chull_bbmin" ].Set( -radius + xlate.X, min + xlate.Y, -radius + xlate.Z );
                chull[ "chull_bbmax" ].Set( radius + xlate.X, max + xlate.Y, radius + xlate.Z );
            }

            // Floor geometry
            Geometry parallelogram = new Geometry( Context );
            parallelogram.PrimitiveCount = 1;
            parallelogram.BoundingBoxProgram = new Program( Context, parrallelPath, "bounds" ); ;
            parallelogram.IntersectionProgram = new Program( Context, parrallelPath, "intersect" ); ;

            Vector3 anchor = new Vector3( -64.0f, 0.01f, -64.0f );
            Vector3 v1 = new Vector3( 128.0f, 0.0f, 0.0f );
            Vector3 v2 = new Vector3( 0.0f, 0.0f, 128.0f );
            Vector3 normal = v2.Cross( v1 );
            normal.Normalize();

            v1 *= 1.0f / ( v1.LengthSqrd() );
            v2 *= 1.0f / ( v2.LengthSqrd() );

            float d = normal.Dot( anchor );
            Vector4 plane = new Vector4( normal, d );

            parallelogram[ "plane" ].Set( ref plane );
            parallelogram[ "v1" ].Set( ref v1 );
            parallelogram[ "v2" ].Set( ref v2 );
            parallelogram[ "anchor" ].Set( ref anchor );

            string boxMtrlName = mTutorial >= 8 ? "box_closest_hit_radiance" : "closest_hit_radiance";
            string floorMtrlName = mTutorial >= 4 ? "floor_closest_hit_radiance" : "closest_hit_radiance";

            Material boxMtrl = new Material( Context );
            boxMtrl.SetSurfaceProgram( 0, new SurfaceProgram( Context, RayHitType.Closest, shaderPath, boxMtrlName ) );
            if ( mTutorial >= 3 )
                boxMtrl.SetSurfaceProgram( 1, new SurfaceProgram( Context, RayHitType.Any, shaderPath, "any_hit_shadow" ) );

            boxMtrl[ "Ka" ].Set( 0.3f, 0.3f, 0.3f );
            boxMtrl[ "Kd" ].Set( 0.6f, 0.7f, 0.8f );
            boxMtrl[ "Ks" ].Set( 0.8f, 0.9f, 0.8f );
            boxMtrl[ "phong_exp" ].Set( 88.0f );
            boxMtrl[ "reflectivity_n" ].Set( 0.2f, 0.2f, 0.2f );

            Material floorMtrl = new Material( Context );
            floorMtrl.SetSurfaceProgram( 0, new SurfaceProgram( Context, RayHitType.Closest, shaderPath, floorMtrlName ) );
            if ( mTutorial >= 3 )
                floorMtrl.SetSurfaceProgram( 1, new SurfaceProgram( Context, RayHitType.Any, shaderPath, "any_hit_shadow" ) );

            floorMtrl[ "Ka" ].Set( 0.3f, 0.3f, 0.1f );
            floorMtrl[ "Kd" ].Set( 194 / 255.0f * .6f, 186 / 255.0f * .6f, 151 / 255.0f * .6f );
            floorMtrl[ "Ks" ].Set( 0.4f, 0.4f, 0.4f );
            floorMtrl[ "reflectivity" ].Set( 0.1f, 0.1f, 0.1f );
            floorMtrl[ "reflectivity_n" ].Set( 0.05f, 0.05f, 0.05f );
            floorMtrl[ "phong_exp" ].Set( 88.0f );
            floorMtrl[ "tile_v0" ].Set( 0.25f, 0, .15f );
            floorMtrl[ "tile_v1" ].Set( -.15f, 0, 0.25f );
            floorMtrl[ "crack_color" ].Set( 0.1f, 0.1f, 0.1f );
            floorMtrl[ "crack_width" ].Set( 0.02f );

            Material glassMtrl = null;
            if ( chull != null )
            {
                glassMtrl = new Material( Context );
                glassMtrl.SetSurfaceProgram( 0, new SurfaceProgram( Context, RayHitType.Closest, shaderPath, "glass_closest_hit_radiance" ) );
                glassMtrl.SetSurfaceProgram( 1, new SurfaceProgram( Context, RayHitType.Any, shaderPath, mTutorial > 9 ? "glass_any_hit_shadow" : "any_hit_shadow" ) );

                Vector3 extinction = new Vector3( .80f, .89f, .75f );
                glassMtrl[ "importance_cutoff" ].Set( 1e-2f );
                glassMtrl[ "cutoff_color" ].Set( 0.34f, 0.55f, 0.85f );
                glassMtrl[ "fresnel_exponent" ].Set( 3.0f );
                glassMtrl[ "fresnel_minimum" ].Set( 0.1f );
                glassMtrl[ "fresnel_maximum" ].Set( 1.0f );
                glassMtrl[ "refraction_index" ].Set( 1.4f );
                glassMtrl[ "refraction_color" ].Set( 1.0f, 1.0f, 1.0f );
                glassMtrl[ "reflection_color" ].Set( 1.0f, 1.0f, 1.0f );
                glassMtrl[ "refraction_maxdepth" ].Set( 100 );
                glassMtrl[ "reflection_maxdepth" ].Set( 100 );
                glassMtrl[ "extinction_constant" ].Set( (float)Math.Log( extinction.X ), (float)Math.Log( extinction.Y ), (float)Math.Log( extinction.Z ) );
                glassMtrl[ "shadow_attenuation" ].Set( 0.4f, 0.7f, 0.4f );
            }

            GeometryInstance boxInst = new GeometryInstance( Context );
            boxInst.Geometry = box;
            boxInst.AddMaterial( boxMtrl );

            GeometryInstance parallelInst = new GeometryInstance( Context );
            parallelInst.Geometry = parallelogram;
            parallelInst.AddMaterial( floorMtrl );
    
            GeometryGroup group = new GeometryGroup( Context );
            group.AddChild( boxInst );
            group.AddChild( parallelInst );
            if ( chull != null )
            {
                GeometryInstance chullInst = new GeometryInstance( Context );
                chullInst.Geometry = chull;
                chullInst.AddMaterial( glassMtrl );
                group.AddChild( chullInst );
            }

            group.Acceleration = new Acceleration( Context, AccelBuilder.Bvh, AccelTraverser.Bvh );

            Context[ "top_object" ].Set( group );
            Context[ "top_shadower" ].Set( group );
        }

        private void LoadEnvMap()
        {
            if ( mTutorial < 5 )
                return;

            string texturePath = Path.GetFullPath( "../../Assets/Textures/CedarCity.png" );

            Bitmap image = new Bitmap( Image.FromFile( texturePath ) );

            if ( image == null )
                return;

            BitmapData imgData = image.LockBits( new Rectangle( 0, 0, image.Width, image.Height ),
                                                ImageLockMode.ReadWrite, image.PixelFormat );


            BufferDesc bufferDesc = new BufferDesc() { Width = (uint)image.Width, Height = (uint)image.Height, Format = Format.UByte4, Type = BufferType.Input };
            Optix.Buffer textureBuffer = new Optix.Buffer( Context, bufferDesc );            

            int stride = imgData.Stride;
            int numChannels = 4;

            unsafe
            {
                byte* src = (byte*)imgData.Scan0.ToPointer();

                BufferStream stream = textureBuffer.Map();
                for ( int h = 0; h < image.Height; h++ )
                {
                    for ( int w = 0; w < image.Width; w++ )
                    {
                        UByte4 color = new UByte4( src[ ( image.Height - h - 1 ) * stride + w * numChannels + 2 ],
                                                   src[ ( image.Height - h - 1 ) * stride + w * numChannels + 1 ],
                                                   src[ ( image.Height - h - 1 ) * stride + w * numChannels + 0 ],
                                                   255 );

                        stream.Write<UByte4>( color );
                    }
                }
                textureBuffer.Unmap();
            }
            image.UnlockBits( imgData );
            
            TextureSampler texture = new TextureSampler( Context, TextureSamplerDesc.GetDefault( WrapMode.Repeat ) );
            texture.SetBuffer( textureBuffer, 0 );

            Context[ "envmap" ].Set( texture );
        }

        private void SetCamera()
        {
            Camera = new Camera();
            Camera.Aspect = (float)Width / (float)Height;
            Camera.Fov = 75.0f;
            Camera.RotationVel = 100.0f;
            Camera.TranslationVel = 500.0f;

            Camera.LookAt( new Vector3( 7.0f, 9.2f, -6.0f ), new Vector3( 0.0f, 4.0f, 0.0f ), Vector3.UnitY );

            CameraUpdate();
        }

        protected override void CameraUpdate()
        {
            Vector3 eye = Camera.Position;
            Vector3 right = Camera.Right;
            Vector3 up = Camera.Up;
            Vector3 look = Camera.Look;

            Context[ "eye" ].Set( ref eye );
            Context[ "U" ].Set( ref right );
            Context[ "V" ].Set( ref up );
            Context[ "W" ].Set( ref look );
        }
    }
}
